%Code for stylized model with EZ preferences in
%'Expectations, Infections, and Economic Activity' by
%M. Eichenbaum, M. Godinho de Matos, F. Lima, S. Rebelo and M. Trabandt

%Fall 2022. mathias.trabandt@gmail.com


clear all;clc; %close all;

%sensitivity options
scaledown=0.23;%lower end of range for sensitivity (relative to baseline value)
scaleup=6; %upper end of range for sensitivity (relative to baseline value)
stepps=99;%steps for sensitivity range


%figure options
ftsize=12;      %Fontsize in plots
lwidth=2;       %Linewidth in plots
ia=1;           %subplot rows
ib=2;           %subplot columns
ymin=0;%0.2;    %y-axis min
ymax=1;%0.6;    %y-axis max

ymin3=0;
ymax3=70;

%parameters
par.betta=0.6; %discount factor
 
par.y1=1;      %income in period 1

par.r=0;        %real interest rate (net)

par.rho=1/1.5;  %Inverse EIS (Albuquerque, Eichenbaum, Luo, Rebelo, benchmark model)
par.alfa=2;     %EZ risk aversion (Albuquerque, Eichenbaum, Luo, Rebelo, benchmark model)

par.same_alfa_rho=0;
if par.same_alfa_rho==1 %set alfa and rho to same value
    par.alfa=par.rho;
end

par.shr_om0_target=0.4365;%share of omega0/(utility of death+bequests), i.e. relative to total payoff from death+bequests
par.payoff_ratio=3.3835; %ratio between payoffs c2/(utility of death+bequests)

par.mu=1-par.rho; %curvature term in bequest utility;

par.delta_c1_target=0.2;  %target value for prob. of dying

par.shr_gam0=0.1;       %baseline target for share of delta due to exogenous constant in delta function
                        %par.delta_c1_target=par.gam0+par.gam1*c1;
                        %1=par.gam0/par.delta_c1_target+par.gam1*c1/par.delta_c1_target;
                        %1=par.shr_gam0+par.gam1*c1/par.delta_c1_target;

par.calibrate_gam0_gam1=1; %switch to calibrate gam0 and gam1; if =0 (below): keep gam0 and gam1 unchanged.
%dont touch here!

par.calibrate_omega0_omega1=1; %switch to calibrate om0 and om1; if =0 (below): keep om0 and om1 unchanged.
%dont touch here!


%fmincon options
opts_fmincon=optimoptions('fmincon','Display','off','ConstraintTolerance',1e-9,'StepTolerance',1e-9,'TolFun',1e-9,...
    'MaxFunctionEvaluations',20000,'UseParallel',false,'FiniteDifferenceStepSize',1e-5,'algorithm','interior-point');

%solve for c1 and c2
guess=[par.y1/2; par.y1/2];
[sol,fval,exitflag]= fmincon(@get_err,guess,[],[],[],[],[],[],@nonlinconstraint,opts_fmincon,par);
if exitflag~=1,error('fmincon did not solve correctly');end

%get c1 and c2
sol_base=sol;
[objective,cc,ceq,c1_base,c2_base,delta_c1_base,delta_c1_prime_base,lam_base,gam0,gam1,omega0,omega1,U1,term]=get_err(sol,par)


par.gam0=gam0;
par.gam1=gam1;
par.omega0=omega0;
par.omega1=omega1;

c1_base=sol(1);
c2_base=sol(2);

%par.calibrate_omega0_omega1=0; %keep omega0 and omega1 unchanged
 
%%%%%%%%%%%%%%%%%%%%%%%%%%%
%comparative statics
%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%
%rho, baseline gam0, gam1
%%%%%%%%%%%%%%%%%%%%%%%%%
par.calibrate_gam0_gam1=0;

par.delta_c1_target_base=par.delta_c1_target;
par.rho_base=par.rho;
par.alfa_base=par.alfa;
par.rho_vec=linspace(par.rho_base*scaledown,par.rho_base*scaleup,stepps);

%for kk=1:1:numel(par.rho_vec)
for kk=numel(par.rho_vec):-1:1
    par.rho=par.rho_vec(kk);
    
    if par.same_alfa_rho==1 %same alfa and rho
        par.alfa=par.rho;
    end
    
    [sol,fval,exitflag]= fmincon(@get_err,sol,[],[],[],[],[],[],@nonlinconstraint,opts_fmincon,par);
    if exitflag~=1,error('fmincon did not solve correctly');end
    [~,~,~,c1,~,~,~,~,~,~,~,~,U1,~]=get_err(sol,par);
       
    c1_vec_rho(kk)=c1;
    U1_vec_rho(kk)=U1;
end
par.rho=par.rho_base;
par.alfa=par.alfa_base;
par.delta_c1_target=par.delta_c1_target_base;

%%%%%%%%%%%%%%%%%%%%%%%%%%
%alfa, baseline gam0, gam1
%%%%%%%%%%%%%%%%%%%%%%%%%%
par.calibrate_gam0_gam1=0;

par.delta_c1_target_base=par.delta_c1_target;
par.rho_base=par.rho;
par.alfa_base=par.alfa;
par.alfa_vec=linspace(par.alfa_base*scaledown,par.alfa_base*scaleup,stepps);

%for kk=1:1:numel(par.alfa_vec)
for kk=numel(par.alfa_vec):-1:1
    par.alfa=par.alfa_vec(kk);
    
    if par.same_alfa_rho==1 %same alfa and rho
        par.rho=par.alfa;
    end
    
    [sol,fval,exitflag]= fmincon(@get_err,sol_base,[],[],[],[],[],[],@nonlinconstraint,opts_fmincon,par);
    if exitflag~=1,error('fmincon did not solve correctly');end  
    [~,~,~,c1,~,~,~,~,~,~,~,~,U1,~]=get_err(sol,par);
       
    c1_vec_alfa(kk)=c1;
    U1_vec_alfa(kk)=U1;
    
    
end
par.alfa=par.alfa_base;
par.rho=par.rho_base;
par.delta_c1_target=par.delta_c1_target_base;






%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%Perturbation: delta=0, i.e. no death prob.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

par.calibrate_gam0_gam1=1;

par.delta_c1_target_perturb2=0;
par.delta_c1_target=par.delta_c1_target_perturb2; %target value for prob. of dying

par.shr_gam0=1;     %baseline target for share of delta due to exogenous constant in delta function
%par.delta_c1_target=par.gam0+par.gam1*c1;
%1=par.gam0/par.delta_c1_target+par.gam1*c1/par.delta_c1_target;
%1=par.shr_gam0+par.gam1*c1/par.delta_c1_target;


%fmincon options
opts_fmincon=optimoptions('fmincon','Display','off','ConstraintTolerance',1e-9,'StepTolerance',1e-9,'TolFun',1e-9,...
    'MaxFunctionEvaluations',20000,'UseParallel',false,'FiniteDifferenceStepSize',1e-5,'algorithm','interior-point');

%solve for c1 and c2
guess=[par.y1/2;par.y1/2];
[sol,fval,exitflag]= fmincon(@get_err,guess,[],[],[],[],[],[],@nonlinconstraint,opts_fmincon,par);
if exitflag~=1,error('fmincon did not solve correctly');end

%get c1 and c2
sol_base=sol;
[objective,cc,ceq,c1_base,c2_base,delta_c1_base,delta_c1_prime_base,lam_base,gam0,gam1,omega0,omega1,U1,term]=get_err(sol,par);
par.gam0=gam0;
par.gam1=gam1;


%%%%%%%%%%%%%%%%%%%%%%%%%
%rho, perturbation 2, no death, i.e. delta=0
%%%%%%%%%%%%%%%%%%%%%%%%%
par.rho_base=par.rho;
par.alfa_base=par.alfa;
par.rho_vec=linspace(par.rho_base*scaledown,par.rho_base*scaleup,stepps);

%for kk=1:1:numel(par.rho_vec)
for kk=numel(par.rho_vec):-1:1
    par.rho=par.rho_vec(kk);
    
    if par.same_alfa_rho==1 %same alfa and rho
        par.alfa=par.rho;
    end
    
    [sol,fval,exitflag]= fmincon(@get_err,sol_base,[],[],[],[],[],[],@nonlinconstraint,opts_fmincon,par);
    if exitflag~=1,error('fmincon did not solve correctly');end    
    
    [~,~,~,c1,~,~,~,~,~,~,~,~,U1,term]=get_err(sol,par);
       
    c1_vec_rho_no_death(kk)=c1;
    U1_vec_rho_no_death(kk)=U1; 
    term_vec_rho_no_death(kk)=term;
    
end
par.rho=par.rho_base;
par.alfa=par.alfa_base;


%%%%%%%%%%%%%%%%%%%%%%%%%
%alfa, perturbation 2, no death, i.e. delta=0
%%%%%%%%%%%%%%%%%%%%%%%%%
par.calibrate_gam0_gam1=0;

par.rho_base=par.rho;
par.alfa_base=par.alfa;
par.alfa_vec=linspace(par.alfa_base*scaledown,par.alfa_base*scaleup,stepps);

%for kk=1:1:numel(par.alfa_vec)
for kk=numel(par.alfa_vec):-1:1
    par.alfa=par.alfa_vec(kk);
    
    if par.same_alfa_rho==1 %same alfa and rho
        par.rho=par.alfa;
    end
    
    [sol,fval,exitflag]= fmincon(@get_err,sol_base,[],[],[],[],[],[],@nonlinconstraint,opts_fmincon,par);
    if exitflag~=1,error('fmincon did not solve correctly');end
    [~,~,~,c1,~,~,~,~,~,~,~,~,U1,term]=get_err(sol,par);
       
    c1_vec_alfa_no_death(kk)=c1;
    U1_vec_alfa_no_death(kk)=U1;
    term_vec_alfa_no_death(kk)=term;
end
par.alfa=par.alfa_base;
par.rho=par.rho_base;






ia=2;ib=2;
%plot results allocations for c1 as functions of alpha and rho
figure;
subplot(ia,ib,1)
plot(par.rho_vec,c1_vec_rho,'LineWidth',lwidth);hold on
plot(par.rho_vec,c1_vec_rho_no_death,'LineWidth',lwidth);hold on
axis([min(par.rho_vec) max(par.rho_vec) ymin ymax]);
vline(par.rho_base);
box off;
legstr2=['\delta>0 (partly endogenous)'];
legstr3=['\delta=',num2str(par.delta_c1_target_perturb2,2),' (no death)'];
leg1=legend(legstr2,legstr3,'Fontsize',ftsize,'Location','northeast');legend box off;

set(leg1,...
    'Position',[0.249696233292831 0.963938529696802 0.50789793438639 0.0264663805436337],...
    'Orientation','horizontal',...
    'FontSize',12);



title('\it  Consumption in Period 1, c_1','interpreter','tex','Fontsize',ftsize+2);
set(gca,'Fontsize',ftsize);
xlabel('\it Inverse EIS, \rho','Fontsize',ftsize+2);

subplot(ia,ib,2)
plot(par.alfa_vec,c1_vec_alfa,'LineWidth',lwidth);hold on;
plot(par.alfa_vec,c1_vec_alfa_no_death,'LineWidth',lwidth);hold on
axis([min(par.alfa_vec) max(par.alfa_vec) ymin ymax]);
vline(par.alfa_base);
box off;
title('\it  Consumption in Period 1, c_1','interpreter','tex','Fontsize',ftsize+2);
set(gca,'Fontsize',ftsize);
xlabel('\it Risk Aversion, \alpha','Fontsize',ftsize+2);

%rho
ybar_rho=U1_vec_rho./term_vec_rho_no_death;
ybar_rho_no_death=U1_vec_rho_no_death./term_vec_rho_no_death;

willingness_rho=(par.y1-ybar_rho)/par.y1*100;
willingness_rho_no_death=(par.y1-ybar_rho_no_death)/par.y1*100;

%alfa
ybar_alfa=U1_vec_alfa./term_vec_alfa_no_death;
ybar_alfa_no_death=U1_vec_alfa_no_death./term_vec_alfa_no_death;

willingness_alfa=(par.y1-ybar_alfa)/par.y1*100;
willingness_alfa_no_death=(par.y1-ybar_alfa_no_death)/par.y1*100;


 


subplot(ia,ib,3)
plot(par.rho_vec,willingness_rho,'LineWidth',lwidth);hold on
plot(par.rho_vec,willingness_rho_no_death,'LineWidth',lwidth);hold on
axis([min(par.rho_vec) max(par.rho_vec) ymin3 ymax3]);
vline(par.rho_base);
box off;
title('\it  Willingness to Pay','interpreter','tex','Fontsize',ftsize+2);
set(gca,'Fontsize',ftsize);
xlabel('\it Inverse EIS, \rho','Fontsize',ftsize+2);
ylabel('\% of income, 100*($y$-$\overline{y}$)/$y$','Fontsize',ftsize+2,'interpreter','latex');

subplot(ia,ib,4)
plot(par.alfa_vec,willingness_alfa,'LineWidth',lwidth);hold on;
plot(par.alfa_vec,willingness_alfa_no_death,'LineWidth',lwidth);hold on
axis([min(par.alfa_vec) max(par.alfa_vec) ymin3 ymax3]);
vline(par.alfa_base);
box off;
title('\it  Willingness to Pay','interpreter','tex','Fontsize',ftsize+2);
set(gca,'Fontsize',ftsize);
xlabel('\it Risk Aversion, \alpha','Fontsize',ftsize+2);
ylabel('\% of income, 100*($y$-$\overline{y}$)/$y$','Fontsize',ftsize+2,'interpreter','latex');


text(-15,-16,'Note: 100*($y$-$\overline{y}$)/$y$ denotes percent of income someone is willing to pay','Fontsize',ftsize+2,'interpreter','latex');
text(-15,-22,'          to remain in $\delta$=0 economy instead of going to $\delta>$0 economy.','Fontsize',ftsize+2,'interpreter','latex');


orient portrait
print -dpdf -fillpage mod_EZstylized

















function [objective,cc,ceq,c1,c2,delta_c1,delta_c1_prime,lam,gam0,gam1,omega0,omega1,U1,term]=get_err(guess,par)

c1=abs(guess(1));
c2=abs(guess(2));
 

if par.calibrate_gam0_gam1
    %back out gam0 and gam1 to hit target shares
    %par.delta_c1_target=par.gam0+par.gam1*c1;
    %1=par.gam0/par.delta_c1_target+par.gam1*c1/par.delta_c1_target;
    gam1=(1-par.shr_gam0)/(c1/par.delta_c1_target);
    gam0=par.shr_gam0*par.delta_c1_target;
    
    par.gam1=gam1;
    par.gam0=gam0;
end

gam0=par.gam0;
gam1=par.gam1;



if par.calibrate_omega0_omega1==1
    %back out om0 and om1 to hit targets
    %shr_om0=c2/(om0+om1*c2^mu)
    %pay_rat=om0/(om0+om1*c2^mu)   
    
    omega1=((1-par.shr_om0_target)*c2/par.payoff_ratio)/c2^par.mu;
    omega0=c2/par.payoff_ratio-omega1*c2^par.mu;
    
    
    par.omega1=omega1;
    par.omega0=omega0;    
    
end

omega0=par.omega0;
omega1=par.omega1;



%calculate delta
delta_c1=par.gam0+par.gam1*c1;
delta_c1_prime=par.gam1;

brace=(  (1-delta_c1)*c2^(1-par.alfa)+delta_c1*(par.omega0+par.omega1*c2^par.mu)^(1-par.alfa)  )^( (1-par.rho)/(1-par.alfa)-1 );
lam=par.betta*brace*( (1-delta_c1 )*c2^(-par.alfa)+delta_c1*(par.omega0+par.omega1*c2^par.mu)^(-par.alfa)*par.mu*par.omega1*c2^(par.mu-1) );

U1=((1-par.betta)*c1^(1-par.rho)+par.betta*((1-delta_c1)*c2^(1-par.alfa)+delta_c1*(par.omega0+par.omega1*c2^par.mu)^(1-par.alfa))^( (1-par.rho)/(1-par.alfa)))^( 1/(1-par.rho) );
term=(1-par.betta)^(1/(1-par.rho))*(1+(par.betta/(1-par.betta))^(1/par.rho))^(par.rho/(1-par.rho));


resid=NaN*zeros(2,1);
resid(1)=-lam*(1+par.r)+(1-par.betta)*c1^(-par.rho)+par.betta/(1-par.alfa)*brace*delta_c1_prime*( (par.omega0+par.omega1*c2^par.mu)^(1-par.alfa)-c2^(1-par.alfa) );
resid(2)=(1+par.r)*(par.y1-c1)-c2;

%resid(3)=c2/(par.omega0+par.omega1*c2^par.mu)-par.payoff_ratio;
%resid(4)=par.omega0/(par.omega0+par.omega1*c2^par.mu)-par.shr_om0_target;


%ceq=0, equality constraint
ceq=resid;

%%cc<=0, inequality constraint.
%cc=NaN*zeros(5,1);
cc(1)=-c1;
cc(2)=-c2;
cc(3)=-lam;
if delta_c1==0
    cc(4)=-10;
else
    cc(4)=-delta_c1;
end
cc(5)=delta_c1-1;
cc(6)=-c2+(par.omega0+par.omega1*c2^par.mu);
%cc(7)=-omega0;
%cc(8)=-omega1;
%cc=[];

%objective to be minimized (trivial)
objective=0;


end


function [cc,ceq] = nonlinconstraint(guess,par)
%get value of nonlinear constraint   must hold
%with equality (ceq)
[~,cc,ceq]=get_err(guess,par);
end


function hhh=vline(x,in1,in2)
% function h=vline(x, linetype, label)
%
% Draws a vertical line on the current axes at the location specified by 'x'.  Optional arguments are
% 'linetype' (default is 'r:') and 'label', which applies a text label to the graph near the line.  The
% label appears in the same color as the line.
%
% The line is held on the current axes, and after plotting the line, the function returns the axes to
% its prior hold state.
%
% The HandleVisibility property of the line object is set to "off", so not only does it not appear on
% legends, but it is not findable by using findobj.  Specifying an output argument causes the function to
% return a handle to the line, so it can be manipulated or deleted.  Also, the HandleVisibility can be
% overridden by setting the root's ShowHiddenHandles property to on.
%
% h = vline(42,'g','The Answer')
%
% returns a handle to a green vertical line on the current axes at x=42, and creates a text object on
% the current axes, close to the line, which reads "The Answer".
%
% vline also supports vector inputs to draw multiple lines at once.  For example,
%
% vline([4 8 12],{'g','r','b'},{'l1','lab2','LABELC'})
%
% draws three lines with the appropriate labels and colors.
%
% By Brandon Kuczenski for Kensington Labs.
% brandon_kuczenski@kensingtonlabs.com
% 8 November 2001

if length(x)>1  % vector input
    for I=1:length(x)
        switch nargin
            case 1
                linetype='r:';
                label='';
            case 2
                if ~iscell(in1)
                    in1={in1};
                end
                if I>length(in1)
                    linetype=in1{end};
                else
                    linetype=in1{I};
                end
                label='';
            case 3
                if ~iscell(in1)
                    in1={in1};
                end
                if ~iscell(in2)
                    in2={in2};
                end
                if I>length(in1)
                    linetype=in1{end};
                else
                    linetype=in1{I};
                end
                if I>length(in2)
                    label=in2{end};
                else
                    label=in2{I};
                end
        end
        h(I)=vline(x(I),linetype,label);
    end
else
    switch nargin
        case 1
            linetype='r:';
            label='';
        case 2
            linetype=in1;
            label='';
        case 3
            linetype=in1;
            label=in2;
    end
    
    
    
    
    g=ishold(gca);
    hold on
    
    y=get(gca,'ylim');
    h=plot([x x],y,linetype,'linewidth',1.5);
    if length(label)
        xx=get(gca,'xlim');
        xrange=xx(2)-xx(1);
        xunit=(x-xx(1))/xrange;
        if xunit<0.8
            text(x+0.01*xrange,y(1)+0.1*(y(2)-y(1)),label,'color',get(h,'color'))
        else
            text(x-.05*xrange,y(1)+0.1*(y(2)-y(1)),label,'color',get(h,'color'))
        end
    end
    
    if g==0
        hold off
    end
    set(h,'tag','vline','handlevisibility','off')
end % else

if nargout
    hhh=h;
end
end